"""
使用 Pinecone 训练和评估 KNN 模型
- 使用 80% 的 MNIST 数据创建 Pinecone 索引
- 使用 20% 的数据测试 k=11 时的准确率
"""

import numpy as np
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from pinecone import Pinecone, ServerlessSpec
import os
import time
from tqdm import tqdm
import logging
from datetime import datetime

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(__name__)

# 初始化 Pinecone
API_KEY = os.getenv("PINECONE_API_KEY", "pcsk_C1E8x_PBoUu1QoydYsUK4K5SWe5EDQvNskevRHiCpc7n8ozSny7Axtf99Ewm2F9LR9F7S")
pc = Pinecone(api_key=API_KEY)

# 索引配置
INDEX_NAME = "mnist-digits-knn"
DIMENSION = 64  # MNIST digits 数据集每个样本 8x8=64 维
BATCH_SIZE = 100  # 每批上传的向量数量

logger.info("=" * 60)
logger.info("开始使用 Pinecone 训练 KNN 模型")
logger.info("=" * 60)

# 加载 MNIST digits 数据集
logger.info("加载 MNIST digits 数据集...")
digits = load_digits()
X = digits.data  # shape: (1797, 64)
y = digits.target  # shape: (1797,)

logger.info(f"数据集总数: {len(X)} 条")

# 划分训练集和测试集（80% 训练，20% 测试）
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

logger.info(f"训练集: {len(X_train)} 条")
logger.info(f"测试集: {len(X_test)} 条")

# 检查索引是否存在
existing_indexes = [index.name for index in pc.list_indexes()]

# 如果索引已存在，先删除
if INDEX_NAME in existing_indexes:
    logger.info(f"删除已存在的索引: {INDEX_NAME}")
    pc.delete_index(INDEX_NAME)
    time.sleep(1)

# 创建新索引
logger.info(f"创建新索引: {INDEX_NAME}")
pc.create_index(
    name=INDEX_NAME,
    dimension=DIMENSION,
    metric="euclidean",
    spec=ServerlessSpec(
        cloud="aws",
        region="us-east-1"
    )
)

# 等待索引创建完成
logger.info("等待索引创建完成...")
while not pc.describe_index(INDEX_NAME).status['ready']:
    time.sleep(1)
logger.info("索引创建完成！")

# 连接到索引
index = pc.Index(INDEX_NAME)

# 上传训练数据到 Pinecone
logger.info("\n开始上传训练数据到 Pinecone...")
logger.info(f"总共需要上传 {len(X_train)} 条数据")

upload_count = 0
with tqdm(total=len(X_train), desc="上传数据", unit="条") as pbar:
    for i in range(0, len(X_train), BATCH_SIZE):
        batch_end = min(i + BATCH_SIZE, len(X_train))
        batch_vectors = []
        
        for j in range(i, batch_end):
            batch_vectors.append({
                "id": f"train_{j}",
                "values": X_train[j].tolist(),
                "metadata": {"label": int(y_train[j])}
            })
        
        # 上传批次
        index.upsert(vectors=batch_vectors)
        upload_count += len(batch_vectors)
        pbar.update(len(batch_vectors))

logger.info(f"成功创建索引并上传了 {upload_count} 条数据")

# 等待索引更新
logger.info("等待索引更新...")
time.sleep(2)

# 验证上传的数据量
stats = index.describe_index_stats()
logger.info(f"索引中的向量数量: {stats['total_vector_count']}")

# 使用测试集评估准确率
logger.info("\n" + "=" * 60)
logger.info("开始评估模型准确率 (k=11)")
logger.info("=" * 60)

k = 11
correct_predictions = 0
total_predictions = len(X_test)

with tqdm(total=total_predictions, desc="测试准确率", unit="条") as pbar:
    for i in range(len(X_test)):
        test_vector = X_test[i].tolist()
        true_label = int(y_test[i])
        
        # 查询最相似的 k 个向量
        query_results = index.query(
            vector=test_vector,
            top_k=k,
            include_metadata=True
        )
        
        # 收集邻居标签
        neighbor_labels = [match['metadata']['label'] for match in query_results['matches']]
        
        # 多数投票
        if neighbor_labels:
            predicted_label = max(set(neighbor_labels), key=neighbor_labels.count)
            
            if predicted_label == true_label:
                correct_predictions += 1
        
        pbar.update(1)

# 计算准确率
accuracy = correct_predictions / total_predictions

logger.info("\n" + "=" * 60)
logger.info("评估结果")
logger.info("=" * 60)
logger.info(f"k 值: {k}")
logger.info(f"测试样本数: {total_predictions}")
logger.info(f"正确预测数: {correct_predictions}")
logger.info(f"使用 Pinecone 的准确率: {accuracy:.4f} ({accuracy*100:.2f}%)")
logger.info("=" * 60)

logger.info("\n训练和评估完成！")
logger.info(f"索引名称: {INDEX_NAME}")
logger.info("你现在可以运行 optimal_knn_webapp_pinecone.py 来启动 Web 应用")
